Enable TBO Support & Fix Accuracy Regressions for Kimi K2.5 - #1369
Enable TBO Support & Fix Accuracy Regressions for Kimi K2.5#1369jpy794 wants to merge 10 commits into
Conversation
There was a problem hiding this comment.
Pull request overview
This PR enables end-to-end Data-Parallel Attention (DPA) + Two-Batch Overlap (TBO) for Kimi K2.5 by fixing DP/TBO micro-batch metadata, strengthening MoE correctness on DP fallback paths, and improving cross-DP prefill admission alignment for TBO’s two-batch requirement.
Changes:
- Propagates per-ubatch per-rank token counts (
ub_tokens_across_dp) through DP sync andForwardContext, and rebuilds ubatch-localDPMetadatainUBatchWrapper. - Fixes MoE correctness/stability under DP fallback + TBO (zero padding rows; scale max token metadata; keepalive tensors across overlapped collectives).
- Updates prefill alignment logic so PrefillDelayer delays until all DP ranks are “alignment-ready” (>=2 local head prefills when TBO is enabled).
Reviewed changes
Copilot reviewed 9 out of 9 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| atom/utils/tbo/ubatching.py | Extends DP sync result to include per-ubatch per-rank token counts for TBO/DP variable-length collectives. |
| atom/utils/tbo/ubatch_wrapper.py | Rebuilds DPMetadata per ubatch using per-ubatch token counts; propagates dp_uniform_decode into ubatch Context. |
| atom/utils/forward_context.py | Adds ub_tokens_across_dp plumbing into ForwardContext / set_forward_context. |
| atom/model_ops/topK.py | Adjusts MORI/all2all gating intended for DPA fallback vs EP mode (but currently has a logic issue). |
| atom/model_ops/moe.py | Zeroes DP all-gather padding rows; scales MoE max token metadata for DP fallback; adds TBO keepalive to prevent premature tensor frees. |
| atom/model_ops/attention_mla.py | Enables persistent MLA for multi-rank DP up to dp_size <= 8. |
| atom/model_engine/scheduler.py | Replaces boolean “prefillable” with counted head-prefill admission and exports both presence + alignment readiness signals. |
| atom/model_engine/prefill_delayer.py | Adds local_alignment_ready and expands MAX-reduce buffer to gate prefill on cross-DP alignment readiness. |
| atom/model_engine/model_runner.py | Threads ub_tokens_across_dp from DP sync into set_forward_context for downstream ubatch/DP metadata. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| and config.enable_expert_parallel | ||
| ) | ||
| if use_mori_all2all: | ||
| return False | ||
| return True |
| return False | ||
| break | ||
|
|
||
| dp_size = config.parallel_config.data_parallel_size |
There was a problem hiding this comment.
This is a duplicate? You can check line 24.
if dp_size > 1 and _has_module("mori") and config.enable_dp_attention: return False
There was a problem hiding this comment.
Thanks for pointing out. That's a rebase error (fixed now). Here we try to enable shared expert fusion for DPA for allgather/reducescatter MoE path (not mori all2all).
| and not self.moe_parallel_config.use_all2all_kernels | ||
| and atom_config.enable_dp_attention | ||
| ): | ||
| moe_max_num_tokens *= self.moe_parallel_config.dp_size |
There was a problem hiding this comment.
I don't understand why we need moe_max_num_tokens *= self.moe_parallel_config.dp_size here.. In all_gahter and model runner, we have padded, * dp_size here will make BS large and kernel bad perf
There was a problem hiding this comment.
Here, we only increase the size of the preallocated internal buffer in FusedMoE, not the actual batch size used in the forward pass. This internal buffer needs to be large enough to accommodate tokens from all DP ranks, so we multiply by dp_size, similar to what we've already done for the all-gather / reduce-scatter buffers.
31ab320 to
426f176
Compare
| @@ -536,9 +538,26 @@ def _can_admit_head_prefill(self) -> bool: | |||
| if num_new_tokens > self.max_num_batched_tokens: | |||
| continue | |||
| if self.block_manager.can_allocate(seq) < 0: | |||
| return False # KV-pressured: definitely cannot prefill | |||
| return True | |||
| return False | |||
| break # KV-pressured: definitely cannot prefill more now. | |||
| if _tbo: | ||
| tbo_switch_to_compute_sync() | ||
| self._hold_tbo_keepalive("ag_output", hidden_states, router_logits) |
| if _tbo: | ||
| tbo_switch_to_compute_sync() | ||
| self._hold_tbo_keepalive("rs_output", final_hidden_states) |
| Mechanism (per scheduler tick): | ||
| 1. Each DP rank reports its local state via cpu all_gather: | ||
| (local_prefillable, watermark_force_allow) | ||
| (local_prefillable, local_alignment_ready, watermark_force_allow) |
There was a problem hiding this comment.
Your changes seem to require >=2 bs for TBO to be ready; does this approach has performance improvement? Or whether it will affect old performance.
There was a problem hiding this comment.
I’ve made this behavior configurable via an environment variable (disabled by default to avoid affecting existing performance).
Below is a Kimi K2.5 performance comparison (conc=128, isl=8k, osl=1k) with ATOM_PREFILL_DELAYER_REQUIRED_PREFILLS=0/2, about 25% throughput gain observed.
ATOM_PREFILL_DELAYER_REQUIRED_PREFILLS=0
============ Serving Benchmark Result ============
Successful requests: 512
Benchmark duration (s): 235.88
Total input tokens: 3775394
Total generated tokens: 473911
Request throughput (req/s): 2.17
Output token throughput (tok/s): 2009.10
Total Token throughput (tok/s): 18014.52
---------------Time to First Token----------------
Mean TTFT (ms): 2994.52
Median TTFT (ms): 786.19
P99 TTFT (ms): 17830.07
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms): 59.09
Median TPOT (ms): 61.84
P99 TPOT (ms): 77.20
---------------Inter-token Latency----------------
Mean ITL (ms): 59.33
Median ITL (ms): 27.85
P99 ITL (ms): 626.02
----------------End-to-end Latency----------------
Mean E2EL (ms): 57912.88
Median E2EL (ms): 58740.95
P99 E2EL (ms): 81474.65
==================================================
ATOM_PREFILL_DELAYER_REQUIRED_PREFILLS=2
============ Serving Benchmark Result ============
Successful requests: 512
Benchmark duration (s): 178.86
Total input tokens: 3775394
Total generated tokens: 473911
Request throughput (req/s): 2.86
Output token throughput (tok/s): 2649.66
Total Token throughput (tok/s): 23758.06
---------------Time to First Token----------------
Mean TTFT (ms): 3435.93
Median TTFT (ms): 1459.46
P99 TTFT (ms): 16975.80
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms): 43.29
Median TPOT (ms): 44.66
P99 TPOT (ms): 54.32
---------------Inter-token Latency----------------
Mean ITL (ms): 43.39
Median ITL (ms): 27.81
P99 ITL (ms): 742.68
----------------End-to-end Latency----------------
Mean E2EL (ms): 43596.09
Median E2EL (ms): 43810.70
P99 E2EL (ms): 60233.29
==================================================
There was a problem hiding this comment.
sure, that's good news. We will test ATOM_PREFILL_DELAYER_REQUIRED_PREFILLS=2 on deepseek v4 and other models if it's indeed effective
|
And could you solve the conflicts and we continue the next step？ |
| # Number of local prefill requests required to allow prefill | ||
| "ATOM_PREFILL_DELAYER_REQUIRED_PREFILLS": lambda: int( | ||
| os.getenv("ATOM_PREFILL_DELAYER_REQUIRED_PREFILLS", "1") | ||
| ), |
| and ( | ||
| self.prefill_delayer is not None | ||
| or self.delay_factor <= 0 | ||
| or self._passed_delay(time.time()) | ||
| ) |
| # A rank counts as "prefillable" for cross-DP alignment only if it | ||
| # can admit a prefill AND has a full batch's worth of waiting tokens. | ||
| # This makes all ranks align on firing dense prefills together | ||
| # instead of straggling partials. |
| # Encoding: | ||
| # slot 0 = local_prefillable (MAX → "any rank prefillable") | ||
| # slot 1 = local_force (MAX → "any rank forces allow") | ||
| # slot 2 = NOT local_prefillable (MAX → "any rank lacks prefill") | ||
| # slot 2 = NOT local_prefill_sufficient | ||
| # (MAX → "any rank lacks required prefill count") | ||
| # Then prefillable_status: |
|
@ZhangLirong-amd Hello, I've rebased to main. I'll add more performance data tomorrow. ATOM config:
|
| ) | ||
|
|
||
| tbo_yield_and_switch_from_compute_to_comm() | ||
| self._hold_tbo_keepalive("ag_source", hidden_states, router_logits) |
There was a problem hiding this comment.
Hi, I have a question, why we need this in all_gather/reduce_scatter with TBO, other models we enabled before didn't meet issues in old logic
There was a problem hiding this comment.
This fixes the use-after-free race for the tensor allocated in stream A and used in stream B. Without this fix the intermediate tensor could be reused by pytorch in the allocating stream before its real use by kernels in the other stream.
You can see the difference in gsm8k (0.9136 vs 0.9515) for Kimi k2.5.
I think other models should have the same race issue, not sure if it's because some minor difference in code path hide this race condition.
Race without tbo_keepalive
==========================
Time ─────────────────────────────────────────────────────────────────────>
Compute stream: produce T ─────────────── drop ref ───── alloc U / reuse T storage
│ │
│ CPU enqueues AG/RS(T) │
▼ ▼
Comm stream: AG/RS reads T ─────────────X
corrupted / UAF
3e64075 to
d2f9b23
Compare
| Mechanism (per scheduler tick): | ||
| 1. Each DP rank reports its local state via cpu all_gather: | ||
| (local_prefillable, watermark_force_allow) | ||
| (local_prefillable, local_prefill_sufficient, watermark_force_allow) |
| and gqa_ratio == 64 | ||
| ) | ||
| use_persistent_mode = dp_size == 1 or requires_persistent_mode | ||
| if envs.ATOM_MLA_PAGE_SIZE > 1: | ||
| use_persistent_mode = False |
| def test_threshold_defaults_to_batch_budget(self): | ||
| cfg = MockConfig(max_num_batched_tokens=32) | ||
| sched = Scheduler(cfg) | ||
| # Threshold is derived from the batch-token budget, not a separate knob. |
| _TBO_KEEPALIVE: dict[tuple[str, int], tuple[torch.Tensor, ...]] = {} | ||
|
|
|
@ZhangLirong-amd Hi, we have refactored the TBO scheduling optimization based on PR #1437 and included a bug fix for an issue in #1437. With this fix, we observed a 13.6% throughput improvement. We have also added more ablation results below and provided the corresponding microbenchmark scripts for the tbo keepalive race condition fix. We hope these additions make the changes easier to evaluate and review. Recently, however, we noticed that PR #1537 reverted #1437 ( It is possible that the bug fix in commit We would greatly appreciate any guidance on how we could adjust or restructure this PR to make it easier to review and merge. We are also happy to split the changes into multiple smaller PRs if that would better align with the upstream development process. Please also feel free to reach out to us for an online discussion if that would be more convenient—we would be happy to walk through the design, implementation, and evaluation results in more detail. MotivationThis PR contains several related TBO changes. If the combined scope is too large for one review, we are happy to discuss splitting it into smaller PRs in the following order:
This PR completes the TBO work in four areas: DP-prefill scheduling, multi-stream correctness, avoiding an unintended padded all-gather path, and Kimi K2.5 enablement. PR #1437 has already upstreamed a cross-DP prefill batching strategy similar to the one proposed here, so this PR does not duplicate that work. Instead, this PR provides the following improvements:
Technical Details1. TBO scheduling optimization
Previously, the scheduler folded two different states into The fix reports both states explicitly:
The cross-DP reduction now delays while any rank is not sufficient, while retaining the existing watermark and timeout escape paths. This aligns the release condition across all DP ranks around batch sufficiency rather than mere request availability. Benchmark configuration: 512 prompts at concurrency 128, random ISL/OSL 8192/1024 with range ratio 0.8,
The 10k rows report the arithmetic mean of three interleaved reruns. In particular, they supersede an earlier 186.52s sufficient-fix run that was affected by run order / machine state. Across the three paired reruns, the sufficient fix reduced duration by 3.33%-4.19%, increased total throughput by 3.45%-4.37%, and reduced P99 TTFT by about 25.5% relative to threshold-only. The batching trace also shows the scheduling effect directly:
The sufficient check eliminates the remaining single-request prefill batches and produces 256 consistently dense two-request batches. 2. TBO correctness bug fix
The fix keeps the previous all-gather/reduce-scatter input and output tensors alive per TBO ubatch and role. They are released only at a later same-role hold, after the ubatch has crossed the synchronization point for the prior communication work. The The end-to-end GSM8K result confirms that this is observable model corruption even though the server starts successfully without the fix. On the full 1,319-example, 5-shot, temperature-0 evaluation at concurrency 128:
Reverting the fix loses 5.7619 percentage points and 76 net correct answers (80 correct-to-wrong and 4 wrong-to-correct). This race is independent of ROCm/aiter#4082. That issue concerns synchronization inside custom collective kernels before callers reuse peer-read input buffers; this PR fixes a separate ATOM/TBO ownership problem where Python tensor references can expire while asynchronous work on another stream still uses their storage. 3. TBO performance bug fix
4. Kimi K2.5 TBO support
Test PlanRun the prefill-delay performance ablation with the following setup and compare both throughput/latency and emitted prefill batch shapes:
Cover no delay, the original 16,384-token threshold, the custom 10k threshold, and the 10k threshold with the sufficient fix. Interleave repeated runs of the last two cases to avoid run-order bias. Run the tensor-lifetime microbenchmark on an otherwise idle GPU: HIP_VISIBLE_DEVICES=0 python tools/multistream_keepalive_race_microbench.pyRun the full GSM8K evaluation with and without |
|
@jpy794 ,sure, please rebase to main and solve the conflict, we will test this branch on dsv4 tbo to make sure its performance. |
* fix(scheduler): gate prefill on full batch to protect decode Hold new prefills until the waiting queue can fill max_num_batched_tokens, else keep decoding. Prevents fast 补发 from firing under-full prefills that preempt decode and drop it out of cudagraph. Tail-escape and pass-budget valves avoid starvation. * style: black format * fix(scheduler): gate dense-batch prefill hold to DP>1 only The prefill dense-batch gate only helps cross-DP rank alignment. Disable it when data_parallel_size<=1 so single-GPU/TP-only runs keep the legacy prefill-first behavior (no added TTFT). --------- Co-authored-by: ZhangLirong-amd <ZhangLirong@amd.com>
d2f9b23 to
56d3dda
Compare
|
@ZhangLirong-amd Hi, I've rebased to main with changes from #1437 included. I'm also glad to help run some dsv4 TBO benchmark to verify the performance, if you could provide some detailed benchmark setup. |
Sure, you can rty |
|
@jpy794 ,hi, seems I meet regression, Total Token throughput (tok/s): 42820.94, in nightly benchmark ,it's about 48000`49000 |
|
For kimi for conc128 + TP4, though haven't tested on other models.
With configs: SemiAnalysisAI/InferenceX@main...amd/zty_test3 
|
|
I could also reproduce the regession in dsv4. Total throughput dropped from 19,345.53 to 18,940.36 with tp8, conc=128, prompts=512. I'm currently investigating the root cause. |
Motivation
Kimi K2.5 inference under Data-Parallel Attention (DPA) combined with Two-Batch Overlap (TBO) exposed several gaps that either crashed the engine or left performance on the table. This PR enables the DPA + TBO path end-to-end: it fixes the fused-MoE fallback and tensor lifetimes in the TBO overlap, aligns cross-DP prefill admission with TBO's two-batch requirement, and extends persistent MLA to multi-rank DP.
Technical Details
Persistent MLA for DP attention (
attention_mla.py): relaxuse_persistent_modefrom "single-rank only" (not (dp_size > 1)) todp_size <= 8, so persistent MLA also runs in the multi-rank DP configuration used by Kimi K2.5.Fix fused MoE on the DPA fallback path (
moe.py,topK.py):dp_size > 1, no MORI all2all), MoE runs afterall_gather_with_padding, so the token dim can grow todp_size ×the per-rank max. Scalemax_num_tokensfor the topK / fused-MoE metadata accordingly to avoid undersized buffers.enable_expert_parallel), so DPA-without-EP correctly falls back instead of assuming all2all.Fix TBO tensor live range (
moe.py): add a per-(role, ubatch)_TBO_KEEPALIVEholder around the all-gather and reduce-scatter comm/compute switches. Under TBO the source/output tensors of in-flight collectives could be freed before the overlapping ubatch waited on the comm; the keepalive defers release to the next same-role hold, which is past the wait point.Two-batch-aware prefill alignment (
scheduler.py,prefill_delayer.py): TBO prefill splitting needs at least two local prefill requests per DP rank. Replace_can_admit_head_prefill(boolean) with_count_admittable_head_prefills(limit)and a_prefill_delayer_readiness()helper that reports both "has any prefill" and "alignment-ready" (>= 2requests when TBO is on,>= 1otherwise).PrefillDelayergains a 4th MAX-reduce slot (local_alignment_ready) so prefill is delayed until every DP rank can launch a full two-batch, not just until one rank has a request.Fix TBO prefill ubatch DP offsets by propagating per-ubatch per-rank token counts through ForwardContext, then rebuilding ubatch-local DPMetadata inside UBatchWrapper. This prevents DP all_gatherv/reduce_scatterv from using full-batch offsets for individual ubatches.
Zero MoE all-gather padding rows before fused-MoE routing/sort/dispatch. Padding rows are later sliced away, but they still participate in fused MoE internals; leaving them uninitialized can introduce NaN/Inf garbage, perturb expert buckets/shared scratch, and corrupt real tokens.
Bugfix Validation
Bad TBO run before MoE padding fix: GSM8K flexible 0.8999, with 125 invalid responses and 238 corrupted outputs.
After fix: GSM8K flexible 0.9742, invalid down to 1, corrupted outputs down to 0.
Perf Benchmark Plan
We tested Kimi K2.5 MXFP4 end-to-end inference on MI355X with ROCm 7.2.3, TP4.
The comparison includes:
Test Results
Numbers in parentheses are throughput/GPU changes relative to the baseline.
At higher concurrency, DPA and TBO show a higher throughput ceiling, with DPA+TBO reaching +15.3% throughput/GPU over the baseline at conc=128.
Submission Checklist